import tensorflow as tf
import numpy as np
import os
from PIL import Image
import datetime
import sys
#IMAGEDATAS_PATH = 'F:\\Dataset\\testimages\\ndid\\1'
PATH_TO_TEST_IMAGES_DIR = 'C:\\Users\\谷雪松\\Documents\\gxs\\datatsets\\VOC2012\\VOCdevkit\\VOC2012\\testJPEGImages' #'test_images'
IMAGEDATAS_PATH = PATH_TO_TEST_IMAGES_DIR #'F:\\Dataset\\VOC0712\\VOCdevkit\\VOC2012\\JPEGImages'
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(2, 4) ]
HASHBASE_PATH = 'object_detection'
ALLOWED_EXTENSIONS = set(['.png', '.jpg', '.jpeg', '.gif'])
X_SIZE = 300
Y_SIZE = 300
ONEHOTNUM = 90
threshold = 0.7
DETECT_PATH = 'object_detection'
sys.path.append("../")
MODEL_NAME = 'models/ssd_mobilenet_v1/inference' # 'ssd_mobilenet_v1_coco_11_06_2017'
PATH_TO_CKPT = os.path.join(MODEL_NAME, 'frozen_inference_graph.pb')

def data_init():
    onehots = []
    imageurls = []
    i = 0
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    detection_graph.as_default()
    sess = tf.Session(graph=detection_graph)
    for parent, dirnames, filenames in os.walk(IMAGEDATAS_PATH):
        for filename in filenames:
            if allowed_files(filename):
                img_path = os.path.join(parent, filename)
                image = Image.open(img_path)
                image_np = load_image_into_numpy_array(image)
                image_np_expanded = np.expand_dims(image_np, axis=0)
                begin1 = datetime.datetime.now()
                scores, hashes, classes = TF_detect(image_np_expanded, sess, detection_graph)
                end1 = datetime.datetime.now()
                print(end1 - begin1)
                # index = classes[np.where(scores > threshold)]
                # index = index.astype(int)
                # out = np.zeros((index.size, ONEHOTNUM))
                #one-hot
                # for row in range(out.shape[0]):
                #     out[row][index[row]-1] = 1
                onehots.append(np.float32(hashes > 0.5))
                # i += 1
                # print(i, index)
                imageurls.append(filename)
    onehot_np = np.array(onehots)
    imageurls_np = np.array(imageurls)
    np.save('onehot_db.npy', onehot_np)
    np.save('urls_db.npy', imageurls_np)
    print("true")
    return onehots, imageurls


def TF_detect(data, sess, detection_graph):
    
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    scores = detection_graph.get_tensor_by_name('detection_scores:0')
    hashes = detection_graph.get_tensor_by_name('detection_hashes:0')
    classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')

    (boxes, scores, hashes, classes, num_detections) = sess.run([boxes, scores, hashes, classes, num_detections],feed_dict={image_tensor: data})
    return scores, hashes, classes


def load_image_into_numpy_array(image):
  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)


def allowed_files(filename):
    return '.' in filename and os.path.splitext(filename)[1] in ALLOWED_EXTENSIONS
if __name__ == '__main__':
    # data_init()
    a = np.load('onehot_db.npy')
    b = np.load('urls_db.npy')
    print(a.shape)
    print(b.size)
